from numpy import ndarray, dtype, floating, float_
from numpy._typing import _64Bit

from factor.BaseFactor import ArrayManagerFactorBase
import talib
import numpy as np
from typing import Union, Literal, Any
from module.dynamic_module.machine_learning.engine import MLEngine


class ArrayManagerFactor(ArrayManagerFactorBase):
    def __init__(self, size, symbol, interval):
        super().__init__(size=size, symbol=symbol, interval=interval)
        # 在因子分析阶段所用到的公共因子（BaseFactor.py内的）
        self.public_factor_ls = ["rsi", "roc", "mom", "atr"]
        # 在因子分析阶段所用到的模型私有因子（BaseFactor.py内的），为空则展示所有私有因子
        self.private_factor_ls = ["sma7", "sma21", "wide_wave", "fast_change", "c_atr", "rise_pro"]
        # # 加载机器学习模型到ml engine
        # self.ml_engine = MLEngine()
        # folder_name = ""
        # self.ml_engine.load_model(folder_name)

    def sma(self, n: int = 5, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        result = np.round(result, decimals=4)
        if array:
            return result
        return result[-1]

    def sma7(self, n: int = 7, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        result = np.round(result, decimals=4)
        if array:
            return result
        return result[-1]

    def sma21(self, n: int = 21, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        result = np.round(result, decimals=4)
        if array:
            return result
        return result[-1]

    def sma60(self, n: int = 60, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        result = np.round(result, decimals=4)
        if array:
            return result
        return result[-1]

    def max21(self, n: int = 21, array: bool = False) -> Union[float, np.ndarray]:
        """
        Max value of period.
        """
        result: np.ndarray = np.zeros(self.size)
        for i in range(self.size):
            if i >= n-1:
                result[i] = np.max(self.close[i-n+1:i+1])
        result = np.round(result, decimals=4)
        if array:
            return result
        return result[-1]

    def min21(self, n: int = 21, array: bool = False) -> Union[float, np.ndarray]:
        """
        Min value of period.
        """
        result: np.ndarray = np.zeros(self.size)
        for i in range(self.size):
            if i >= n-1:
                result[i] = np.min(self.close[i-n+1:i+1])
        result = np.round(result, decimals=4)
        if array:
            return result
        return result[-1]

    def dd(self, n: int = 0, array: bool = False) -> Union[float, np.ndarray]:
        """
        自定义因子dd，因子计算方法为后一个K线的close减去前一个K线的open，差为因子值
        """
        # 初始化数据维度与self.close一致的result数据
        result: np.ndarray = np.zeros(self.size)
        # 覆盖result中对应位置的数据为因子值
        result[1:] = self.close[1:] - self.open[:1]
        result = np.round(result, decimals=4)
        if array:
            return result
        return result[-1]

    def c_atr(self, n: int = 10, array=False) -> Union[float, np.array]:
        atr = self.atr(n, array=True)
        c_atr = np.round(np.power(atr*100/self.close, 1/4), 4)
        result = c_atr
        if array:
            return result
        else:
            return result[-1]

    def n_write_soldiers(self, n=3, array=False):
        result = np.zeros(self.size)
        for i in range(self.size):
            if i < n - 1:
                pass
            else:
                num = 0
                for j in range(n):
                    if i == n - 1:
                        continue
                    now_index = i-j
                    # 连续三根蜡烛都是上涨的
                    up_cond = self.close[now_index] > self.open[now_index]
                    # 每根蜡烛的收盘价都高于前一根的收盘价
                    above_cond = self.close[now_index] > self.close[now_index-1]
                    # 每根蜡烛的开盘价都在前一根蜡烛的实体内（可选）
                    inside_cond = self.open[now_index-1] < self.open[now_index] < self.close[now_index-1]

                    if up_cond and above_cond and inside_cond:
                        num += 1
                        if num >= n:
                            result[i] = 1

        if array:
            return result
        else:
            return result[-1]



